import os
import pandas as pd
from collections import Counter
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
import joblib

# 显示所有列
pd.set_option('display.max_columns', None)

# 加载泰坦尼克数据集
trainic_data = pd.read_csv("data/train.csv")
# 加载测试数据集
test_data = pd.read_csv("data/test.csv")
train_data = trainic_data.append(test_data, ignore_index=True)

# 年龄（Age）和船票价格（Fare）缺失值用平均数代替
train_data['Age'] = train_data['Age'].fillna(train_data['Age'].mean())
train_data['Fare'] = train_data['Fare'].fillna(train_data['Fare'].mean())
# 客舱号（Cabin）缺失值用Unknown代替
train_data['Cabin'] = train_data['Cabin'].fillna('Unknown')
# 登船口（Embarked）缺失值用出现最多的S代替
train_data['Embarked'] = train_data['Embarked'].fillna('S')

sex_map = {'male': 1, 'female': 0}
train_data['Sex'] = train_data['Sex'].map(sex_map)

# 登船口和客舱等级数据one-host
embarkedDf = pd.DataFrame()
embarkedDf = pd.get_dummies(train_data['Embarked'], prefix='Embarked')
train_data = pd.concat([train_data, embarkedDf], axis=1)
train_data.drop('Embarked', axis=1, inplace=True)

pclassDf = pd.DataFrame()
pclassDf = pd.get_dummies(train_data['Pclass'], prefix='Pclass')
train_data = pd.concat([train_data, pclassDf], axis=1)
train_data.drop('Pclass', axis=1, inplace=True)


# 提取姓名头衔
def getNameTitle(name):
    str1 = name.split(",")[1]
    str2 = str1.split(".")[0]
    str3 = str2.strip()
    return str3


nameTitle_map = {
    'Capt': 'Officer',
    'Col': 'Officer',
    'Major': 'Officer',
    'Jonkheer': 'Royalty',
    'Don': 'Royalty',
    'Sir': 'Royalty',
    'Dr': 'Officer',
    'Rev': 'Officer',
    'the Countess': 'Royalty',
    'Dona': 'Royalty',
    'Mme': 'Mrs',
    'Mlle': 'Miss',
    'Ms': 'Mrs',
    'Mr': 'Mr',
    'Mrs': 'Mrs',
    'Miss': 'Miss',
    'Master': 'Master',
    'Lady': 'Royalty',
}

nameTitleDf = pd.DataFrame()
nameTitleDf['Title'] = train_data['Name'].map(getNameTitle)
nameTitleDf['Title'] = nameTitleDf['Title'].map(nameTitle_map)
nameTitleDf = pd.get_dummies(nameTitleDf['Title'])
train_data = pd.concat([train_data, nameTitleDf], axis=1)
train_data.drop('Name', axis=1, inplace=True)

# 提取客舱信息（Cabin）
cabinDf = pd.DataFrame()
train_data['Cabin'] = train_data['Cabin'].map(lambda c: c[0])
cabinDf = pd.get_dummies(train_data['Cabin'], prefix='Cabin')
train_data = pd.concat([train_data, cabinDf], axis=1)
train_data.drop('Cabin', axis=1, inplace=True)

# 家庭信息，按1,5划分小家庭、中等家庭和大家庭
familyDf = pd.DataFrame()
familyDf['Familysize'] = train_data['SibSp'] + train_data['Parch'] + 1
familyDf['Familysize_Single'] = familyDf['Familysize'].map(lambda s: 1 if s <= 1 else 0)
familyDf['Familysize_Small'] = familyDf['Familysize'].map(lambda s: 1 if 1 < s < 5 else 0)
familyDf['Familysize_Large'] = familyDf['Familysize'].map(lambda s: 1 if 5 <= s else 0)
train_data = pd.concat([train_data, familyDf], axis=1)

# 年龄信息，按6、18、40、60分为幼、少、青、中、老年
ageDf = pd.DataFrame()
ageDf['Child'] = train_data['Age'].map(lambda a: 1 if 0 < a <= 6 else 0)
ageDf['Teenager'] = train_data['Age'].map(lambda a: 1 if 6 < a <= 18 else 0)
ageDf['Youth'] = train_data['Age'].map(lambda a: 1 if 18 < a <= 40 else 0)
ageDf['Middle_aged'] = train_data['Age'].map(lambda a: 1 if 40 < a <= 60 else 0)
ageDf['Older'] = train_data['Age'].map(lambda a: 1 if 60 < a else 0)
train_data = pd.concat([train_data, ageDf], axis=1)
train_data.drop('Age', axis=1, inplace=True)

# 查看各项与生还的相关关系
corrDf = train_data.corr()
correlation = corrDf['Survived'].sort_values(ascending=False)
# print(correlation)

train_dataX = pd.concat([
    train_data['Sex'],  # 性别
    pclassDf,  # 客舱等级
    familyDf,  # 家庭情况
    nameTitleDf,  # 头衔
    ageDf,  # 年龄
    train_data['Fare'],  # 票价
    cabinDf,  # 客舱号
    embarkedDf  # 登船地点
], axis=1)
# print(train_dataX)


# 训练模型
sourceRow = 891
source_X = train_dataX.loc[0:sourceRow - 1, :]
source_Y = train_data.loc[0:sourceRow - 1, 'Survived']
pred_X = train_dataX.loc[sourceRow:, :]

# 随机抽取样本检测
train_X, test_X, train_Y, test_Y = train_test_split(source_X, source_Y, train_size=0.8)

# 逻辑回归模型
model = LogisticRegression(random_state=0, solver='lbfgs', max_iter=3000, multi_class='multinomial')
model.fit(train_X, train_Y)
# model.score(test_X, test_Y) 检测模型准确率
pred_y = model.predict(pred_X)  # 结果预测
pred_y = pred_y.astype(int)
pred_y_detail = model.predict_proba(pred_X)  # 比例预测
survivedTrue = []
survivedFalse = []
for result in pred_y_detail:
    survivedTrue.append(result[1])
    survivedFalse.append(result[0])
passenger_id = train_data.loc[sourceRow:, 'PassengerId']
predDf = pd.DataFrame(
    {'PassengerId': passenger_id, 'Survived': pred_y, 'SurvivedTrue': survivedTrue, 'SurvivedFalse': survivedFalse})

# 保存结果集
predDf.to_csv('data/titanic_pred.csv', index=False)

# 保存模型
joblib.dump(model, 'modle/titanic.modle')
